#!/usr/bin/env python3
import sys, os
import argparse
import datetime
import yaml # pip install pyyaml
import logging
logging.basicConfig(stream=sys.stderr, level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')

if __name__ == '__main__':
    logging.info('Host Runner Started')

    parser = argparse.ArgumentParser()
    parser.add_argument('--stage', type=int, default=0, help='set stage variable in benchmark.sh')
    parser.add_argument('--max_num_utts', type=int, default=10000, help='max number of utts to test in a test set')
    parser.add_argument('request_file_path', type=str, help='request.yaml')
    args = parser.parse_args()

    model_zoo   = os.path.abspath('models');   assert os.path.isdir(model_zoo)
    dataset_zoo = os.path.abspath('datasets'); assert os.path.isdir(dataset_zoo)

    logging.info('Processing benchmark request ...')
    with open(args.request_file_path, 'r', encoding = 'utf-8') as request_file:
        request = yaml.safe_load(request_file)
        model_id = request['model']
        test_sets = []
        for x in request['test_set']:
            test_sets.append(x)
    assert(model_id)
    assert(test_sets)

    model_dir = os.path.join(model_zoo, model_id)

    # download model
    if not os.path.isdir(model_dir):
        pull_model_cmd = F'ops/pull model {model_id}'
        logging.info(F'Please pull model-image to your local machine via: {pull_model_cmd}')
        sys.exit(-1)
    assert os.path.isdir(model_dir)

    # check model
    model_info_path = os.path.join(model_dir, 'model.yaml')
    assert os.path.isfile(model_info_path)
    sbi_bin_path = os.path.join(model_dir, 'SBI')
    assert os.path.isfile(sbi_bin_path)
    docker_context = os.path.join(model_dir, 'docker')
    assert os.path.isdir(docker_context)
    docker_file = os.path.join(docker_context, 'Dockerfile')
    assert os.path.isfile(docker_file)

    with open(model_info_path, 'r', encoding = 'utf-8') as model_info_file:
        model_info = yaml.safe_load(model_info_file)
        task = model_info['task']
        language = model_info['language']
    assert task == 'ASR'
    assert language

    # build docker image
    docker_image = F'speechio/leaderboard:{model_id}'
    docker_build_cmd = (
        F'docker build -f {docker_file} -t {docker_image} '
        #F'--build-arg USER_ID=$(id -u) --build-arg GROUP_ID=$(id -g) '
        F'{docker_context}'
    )
    logging.info(F'Building docker image: {docker_build_cmd}')
    os.system(docker_build_cmd)

    # run benchmarking docker
    testset_ids = ' '.join(test_sets)
    leaderboard = os.getcwd()
    LEADERBOARD = '/app/speechio/leaderboard'
    docker_run_cmd = (
        F'docker run '
        F'--user "$(id -u):$(id -g)" '
        F'-e LEADERBOARD={LEADERBOARD} '
        F'-v {leaderboard}:{LEADERBOARD} '
        F'{docker_image} '
        F'{LEADERBOARD}/utils/benchmark.sh --stage {args.stage} --max-num-utts {args.max_num_utts} {model_id} "{testset_ids}"'
    )
    logging.info(F'Pulling up benchmarking container: {docker_run_cmd}')
    os.system(docker_run_cmd)

    logging.info('Host Runner Done')
